'''
使用没有进行数据增强的初始数据集进行数据集生成
保存在origin_data数据文件夹下
用于训练未采用过采样的模型
'''

import os
import csv
import pydicom
import shutil
import numpy as np
import pandas as pd

# 得到指定Instance Number的dicom图像名字
def getName(origin_path, SeriesDescription, InstanceNumber):
    names = os.listdir(origin_path)

    for name in names:
        slice = pydicom.read_file(origin_path + '/' + name)
        if slice.SeriesDescription == SeriesDescription and slice.InstanceNumber == InstanceNumber:
            return name

    print('------------------------------------------', origin_path.split('/')[-1], '-----------------------------')

    return 'no'


# 读取labels.csv的CT和PET信息，生成 all_data1.csv
def save_info():

    num = 1

    columns = ['patientID', 'z', 'x', 'y', 'r', 'cancer_type', 'ct_count', 'pet_count', 'ct_size', 'pet_size',
               'ct_x_spacing', 'ct_y_spacing', 'ct_z_spacing', 'pet_x_spacing', 'pet_y_spacing', 'pet_z_spacing',
               'ct_intercept', 'pet_intercept', 'ct_slope', 'pet_slope', 'patientWeight', 'patientSex', 'patientAge',
               'patientSize', 'TotalDose', 't0', 't1', 'HalfLife', 'DecayFactor', 'CT_SeriesDescription',
               'PET_SeriesDescription', 'origin_dir', 'CT_origin_path', 'PET_origin_path']

    myfile = open('D:/lung_cancer/data/data_augmentation/labels.csv')

    labels = []
    lines = csv.reader(myfile)
    for line in lines:
        labels.append(line)

    infos = []
    for line in labels[1:]:

        path = 'H:/'+line[9]
        names = os.listdir(path)
        ct_slices = []
        pet_slices = []

        for name in names:
            slice = pydicom.read_file(path + '/' + name)
            if slice.SeriesDescription == str(line[7]):
                ct_slices.append(slice)
            elif slice.SeriesDescription == str(line[8]):
                pet_slices.append(slice)
        ct_slices.sort(key=lambda x: int(x.InstanceNumber))  # 默认升序排列
        pet_slices.sort(key=lambda y: int(y.InstanceNumber))  # 默认升序排列

        ct_count = len(ct_slices)
        pet_count = len(pet_slices)


        ct_z_spacing = np.abs(ct_slices[0].SliceLocation - ct_slices[1].SliceLocation)
        pet_z_spacing = np.abs(pet_slices[0].SliceLocation - pet_slices[1].SliceLocation)

        patientID = int(line[1])
        z = int(line[2])
        x = int(line[3])
        y = int(line[4])
        r = int(line[5])
        cancer_type = int(line[6])

        CT_SeriesDescription = line[7]
        PET_SeriesDescription = line[8]

        # 得到原始地址
        CT_origin_path = getName(path, CT_SeriesDescription, z)
        PET_origin_path = getName(path, PET_SeriesDescription, z)

        # 方法一：用附近的代替
        # while CT_origin_path == 'no' or PET_origin_path == 'no':
        #     z = z+1
        #     CT_origin_path = getName(path, CT_SeriesDescription, z)
        #     PET_origin_path = getName(path, PET_SeriesDescription, z)

        # 方法二：直接跳过
        if CT_origin_path == 'no' or PET_origin_path == 'no':
            continue

        CT_origin_path = line[9]+'/'+getName(path, CT_SeriesDescription, z)
        PET_origin_path = line[9]+'/'+getName(path, PET_SeriesDescription, z)


        # 可能因为ct和pet的Instance Number不一定从0开始，所以不能直接用上文的索引得到ct和pet
        ct = pydicom.read_file('H:/'+CT_origin_path)
        pet = pydicom.read_file('H:/'+PET_origin_path)

        ct_size = ct.Rows
        pet_size = pet.Rows
        ct_x_spacing = ct.PixelSpacing[0]
        ct_y_spacing = ct.PixelSpacing[1]
        pet_x_spacing = pet.PixelSpacing[0]
        pet_y_spacing = pet.PixelSpacing[1]
        ct_intercept = ct.RescaleIntercept
        pet_intercept = pet.RescaleIntercept
        ct_slope = ct.RescaleSlope
        pet_slope = pet.RescaleSlope

        patientAge = int(ct[0x10, 0x1010].value[:-1])
        patientWeight = float(pet[0x10, 0x1030].value)
        patientSize = float(pet[0x10, 0x1020].value) * 100
        Sex = pet[0x10, 0x40].value
        patientSex = 0  # 男性为0， 女性为1
        if (Sex == 'F'):
            patientSex = 1
        DecayFactor = pet.DecayFactor

        a = pet.RadiopharmaceuticalInformationSequence
        TotalDose = a[0].RadionuclideTotalDose
        HalfLife = a[0].RadionuclideHalfLife
        t0 = a[0].RadiopharmaceuticalStartTime
        t1 = pet.AcquisitionTime


        origin_dir = line[9]



        infos.append([patientID, z, x, y, r, cancer_type, ct_count, pet_count, ct_size, pet_size, ct_x_spacing,
                      ct_y_spacing, ct_z_spacing, pet_x_spacing, pet_y_spacing, pet_z_spacing, ct_intercept,
                      pet_intercept, ct_slope, pet_slope, patientWeight, patientSex, patientAge,patientSize,
                      TotalDose, t0, t1, HalfLife, DecayFactor, CT_SeriesDescription, PET_SeriesDescription, origin_dir,
                      CT_origin_path, PET_origin_path])
        print(num, '--->', patientID, 'is ok')
        num = num+1
    df = pd.DataFrame(infos, columns=columns)
    # df.to_csv('D:/lung_cancer/data/data_augmentation/all_data1.csv', index=False)
    df.to_csv('D:/lung_cancer/data/data_augmentation/all_data1.csv', index=False)

# 生成的身高有些问题，重新格式化身高, 生成新的all_data1.csv
def standardPatientSize():
    myfile = open('D:/lung_cancer/data/data_augmentation/all_data1.csv')


    labels = []
    lines = csv.reader(myfile)
    for line in lines:
        labels.append(line)

    infos = []
    for line in labels[1:]:
        patientSize = int(float(line[23]))
        print(patientSize)
        if patientSize > 999:
            print(line[0])
            line[23] = int(patientSize/100)
        infos.append(line)

    df = pd.DataFrame(infos, columns=labels[0])
    df.to_csv('D:/lung_cancer/data/data_augmentation/all_data1.csv', index=False)



# 利用all_data1.csv中的文件地址，把要用到的CT和PET单独保存
def extract_origin_CT_PET():
    data = pd.read_csv('D:/lung_cancer/data/data_augmentation/all_data1.csv')

    CT_path = data['CT_origin_path']
    PET_path = data['PET_origin_path']
    patientID = data['patientID']
    print(patientID[len(patientID)-1])

    for i in range(len(patientID)):
        print(i, '----->', patientID[i])
        CT_save_path = 'D:/lung_cancer/data/data_augmentation/origin_data/'+str(patientID[i])+'/CT/'
        CT_origin_path = 'H:/'+CT_path[i]

        PET_save_path = 'D:/lung_cancer/data/data_augmentation/origin_data/'+str(patientID[i])+'/PET/'
        PET_origin_path = 'H:/' + PET_path[i]

        if (not os.path.exists(CT_save_path)):
            os.makedirs(CT_save_path)
        shutil.copy(CT_origin_path, CT_save_path)

        if (not os.path.exists(PET_save_path)):
            os.makedirs(PET_save_path)
        shutil.copy(PET_origin_path, PET_save_path)

# 经过cut_pet.py-->cut_ct.py-->calculate_suv.py生成all_data4.csv
# 将肺部长宽等信息加进来生成all_data5.csv
def get_lungHW():
    # 之前提取的含有肺部长宽的csv文件
    data_path = 'D:/lung_cancer/data/all_data7.csv'
    data = []
    f = csv.reader(open(data_path, 'r'))
    for i in f:
        data.append(i)

    # 要存入的文件
    save_data_path = 'D:/lung_cancer/data/data_augmentation/all_data4.csv'
    save_data = []
    f2 = csv.reader(open(save_data_path, 'r'))
    for item in f2:
        save_data.append(item)

    new_data = []

    for i in range(1, len(save_data)):
        line = list(save_data[i])
        line.append(data[i][-8])
        line.append(data[i][-7])
        line.append(data[i][-6])
        line.append(data[i][-5])
        line.append(data[i][-4])
        line.append(data[i][-3])
        line.append(data[i][-2])
        line.append(data[i][-1])
        new_data.append(line)


    save_data[0].append('center_x')
    save_data[0].append('center_y')
    save_data[0].append('lungW')
    save_data[0].append('lungH')
    save_data[0].append('newx')
    save_data[0].append('newy')
    save_data[0].append('cx')
    save_data[0].append('cy')

    df = pd.DataFrame(new_data, columns=save_data[0])
    df.to_csv('D:/lung_cancer/data/data_augmentation/all_data5.csv', index=False)


# 以医生预测样本为测试集，划分训练集和测试集
# split_data.py

if __name__ == '__main__':
    # save_info()
    # standardPatientSize()
    get_lungHW()